#!/usr/bin/env python
import rospy
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from geometry_msgs.msg import PoseStamped
from nav_msgs.msg import Odometry
from collections import deque

# 第一帧动捕位置
x = 0.0
y = 0.0
z = 0.0
flag = True

# 全局变量存储位姿数据
max_huav1_length = 400  # 保留的最大历史数据长度
max_odom_length = 100
uav1_pose_history = deque(maxlen = max_huav1_length)  # 使用 deque 存储 UAV1 的历史位姿
baton_odom_history = deque(maxlen = max_odom_length)  # 使用 deque 存储 Baton 的历史位姿

# 回调函数：订阅 /vrpn_client_node/uav1/pose 话题
def uav1_pose_callback(msg):
    global flag, x, y, z
    if flag:
        flag = False
        x = msg.pose.position.x
        y = msg.pose.position.y
        z = msg.pose.position.z
        print("x: ", x, " y: ", y, " z: ", z)

    position = msg.pose.position
    uav1_pose_history.append((position.x, position.y, position.z))  # 存储历史位姿

# 回调函数：订阅 /baton/stereo3/odometry 话题
def baton_odom_callback(msg):
    global x, y, z
    if not flag:  # 确保第一帧动捕位置已经记录
        position = msg.pose.pose.position
        # 计算相对于第一帧的偏移量
        baton_odom_history.append((position.x + x, position.y + y, position.z + z))  # 存储历史位姿

# 动态调整坐标轴范围
def update_axes(ax, uav1_history, baton_history):
    if not uav1_history and not baton_history:
        return  # 如果没有数据，不调整坐标轴
    
    # 获取所有历史数据的范围
    all_points = list(uav1_history) + list(baton_history)
    x_vals = [p[0] for p in all_points]
    y_vals = [p[1] for p in all_points]
    z_vals = [p[2] for p in all_points]
    
    # 设置坐标轴范围
    ax.set_xlim(min(x_vals) - 1, max(x_vals) + 1)
    ax.set_ylim(min(y_vals) - 1, max(y_vals) + 1)
    ax.set_zlim(min(z_vals) - 1, max(z_vals) + 1)

# 主函数
if __name__ == '__main__':
    # 初始化 ROS 节点
    rospy.init_node('realtime_pose_plotter')

    # 订阅话题
    rospy.Subscriber('/vrpn_client_node/vibo_1/pose', PoseStamped, uav1_pose_callback)
    rospy.Subscriber('/baton/stereo3/odometry', Odometry, baton_odom_callback)

    # 创建实时绘图
    plt.ion()  # 开启交互模式
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.set_xlabel('X Axis')
    ax.set_ylabel('Y Axis')
    ax.set_zlabel('Z Axis')
    ax.set_title('Real-time 3D Pose Visualization with History')
    ax.set_xlim(-3,3)
    ax.set_ylim(-3,3)
    ax.set_zlim(-1,1.5)

    # 初始化绘图数据
    uav1_line, = ax.plot([], [], [], c='r', marker='o', label='UAV1 Pose', linestyle='-', markersize=3)
    baton_line, = ax.plot([], [], [], c='b', marker='^', label='Baton Odometry', linestyle='-', markersize=3)
    ax.legend()

    # 实时更新绘图
    try:
        while not rospy.is_shutdown():
            if uav1_pose_history or baton_odom_history:
                # 提取 UAV1 和 Baton 的历史位姿
                uav1_x = [p[0] for p in uav1_pose_history]
                uav1_y = [p[1] for p in uav1_pose_history]
                uav1_z = [p[2] for p in uav1_pose_history]
                
                baton_x = [p[0] for p in baton_odom_history]
                baton_y = [p[1] for p in baton_odom_history]
                baton_z = [p[2] for p in baton_odom_history]
                
                # 更新 UAV1 轨迹
                uav1_line.set_data(uav1_x, uav1_y)
                uav1_line.set_3d_properties(uav1_z)
                
                # 更新 Baton 轨迹
                baton_line.set_data(baton_x, baton_y)
                baton_line.set_3d_properties(baton_z)
                
                # 动态调整坐标轴范围
                # update_axes(ax, uav1_pose_history, baton_odom_history)

                # 重绘图
                plt.draw()
                plt.pause(0.01)  # 稍微暂停以更新图形
            
            rospy.sleep(0.1)  # 控制循环频率
    except KeyboardInterrupt:
        print("程序被用户中断")
    finally:
        rospy.signal_shutdown("程序结束")  # 确保 ROS 节点关闭
        plt.close()  # 关闭图形窗口